﻿/* KeeChallenge--Provides Yubikey challenge-response capability to Keepass
*  Copyright (C) 2014  Ben Rush
*  
*  This program is free software; you can redistribute it and/or
*  modify it under the terms of the GNU General Public License
*  as published by the Free Software Foundation; either version 2
*  of the License, or (at your option) any later version.
*  
*  This program is distributed in the hope that it will be useful,
*  but WITHOUT ANY WARRANTY; without even the implied warranty of
*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
*  GNU General Public License for more details.
*  
*  You should have received a copy of the GNU General Public License
*  along with this program; if not, write to the Free Software
*  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
*/

using System;
using System.IO;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Security.Cryptography;
using System.Diagnostics;
using System.Xml;

using KeePassLib.Keys;
using KeePassLib.Utility;
using KeePassLib.Cryptography;
using KeePassLib.Cryptography.Cipher;
using KeePassLib.Serialization;

using keepass2android;
using keepass2android.Io;

namespace KeeChallenge
{
  public sealed class KeeChallengeProv
  {
    private const string m_name = "Yubikey challenge-response";

    public static string Name { get { return m_name; } }

    public const int keyLenBytes = 20;
    public const int challengeLenBytes = 64;
    public const int responseLenBytes = 20;
    public const int secretLenBytes = 20;

    //If variable length challenges are enabled, a 63 byte challenge is sent instead.
    //See GenerateChallenge() and http://forum.yubico.com/viewtopic.php?f=16&t=1078
    //This field is automatically set by calling GetSecret(). However, when creating 
    //a new database it will need to be set manually based on the user's yubikey settings
    public bool LT64
    {
      get;
      set;
    }

    public KeeChallengeProv()
    {
      LT64 = false;
    }

    private byte[] GenerateChallenge()
    {
      byte[] chal = CryptoRandom.Instance.GetRandomBytes(challengeLenBytes);
      if (LT64)
      {
        chal[challengeLenBytes - 2] = (byte)~chal[challengeLenBytes - 1];
      }

      return chal;
    }

    private byte[] GenerateResponse(byte[] challenge, byte[] key)
    {
      HMACSHA1 hmac = new HMACSHA1(key);

      if (LT64)
        challenge = challenge.Take(challengeLenBytes - 1).ToArray();

      byte[] resp = hmac.ComputeHash(challenge);
      hmac.Clear();
      return resp;
    }

    /// <summary>
    /// A method for generating encrypted ChallengeInfo to be saved. For security, this method should
    /// be called every time you get a successful challenge-response pair from the Yubikey. Failure to
    /// do so will permit password re-use attacks. 
    /// </summary>
    /// <param name="secret">The un-encrypted secret</param>
    /// <returns>A fully populated ChallengeInfo object ready to be saved</returns>
    public ChallengeInfo Encrypt(byte[] secret)
    {
      //generate a random challenge for use next time
      byte[] challenge = GenerateChallenge();

      //generate the expected HMAC-SHA1 response for the challenge based on the secret
      byte[] resp = GenerateResponse(challenge, secret);

      //use the response to encrypt the secret
      SHA256 sha = SHA256Managed.Create();
      byte[] key = sha.ComputeHash(resp); // get a 256 bit key from the 160 bit hmac response
      byte[] secretHash = sha.ComputeHash(secret);

      StandardAesEngine aes = new StandardAesEngine();
      const uint aesIVLenBytes = 16;
      byte[] IV = CryptoRandom.Instance.GetRandomBytes(aesIVLenBytes);
      byte[] encrypted;

      using (MemoryStream msEncrypt = new MemoryStream())
      {
        using (CryptoStream csEncrypt = (CryptoStream)aes.EncryptStream(msEncrypt, key, IV))
        {
          csEncrypt.Write(secret, 0, secret.Length);
          csEncrypt.Close();
        }

        encrypted = msEncrypt.ToArray();
        msEncrypt.Close();
      }

      ChallengeInfo inf = new ChallengeInfo(encrypted, IV, challenge, secretHash, LT64);

      sha.Clear();

      return inf;
    }

    private bool DecryptSecret(byte[] yubiResp, ChallengeInfo inf, out byte[] secret)
    {
      secret = new byte[keyLenBytes];

      if (inf.IV == null) return false;
      if (inf.Verification == null) return false;

      //use the response to decrypt the secret
      SHA256 sha = SHA256Managed.Create();
      byte[] key = sha.ComputeHash(yubiResp); // get a 256 bit key from the 160 bit hmac response

      StandardAesEngine aes = new StandardAesEngine();

      using (MemoryStream msDecrypt = new MemoryStream(inf.EncryptedSecret))
      {
        using (CryptoStream csDecrypt = (CryptoStream)aes.DecryptStream(msDecrypt, key, inf.IV))
        {
          //read the secret from the stream
          int totalBytesRead = 0;

          int bytesRead = csDecrypt.Read(secret, totalBytesRead, secret.Length - totalBytesRead);
          while (bytesRead > 0 && totalBytesRead < secret.Length)
          {
            totalBytesRead += bytesRead;
            bytesRead = csDecrypt.Read(secret, totalBytesRead, secret.Length - totalBytesRead);
          }

          csDecrypt.Close();
        }
        msDecrypt.Close();
      }

      byte[] secretHash = sha.ComputeHash(secret);
      for (int i = 0; i < secretHash.Length; i++)
      {
        if (secretHash[i] != inf.Verification[i])
        {
          //wrong response
          Array.Clear(secret, 0, secret.Length);
          return false;
        }
      }

      //return the secret
      sha.Clear();
      return true;
    }

    /// <summary>
    /// The primary access point for challenge-response utility functions. Accepts a pre-populated ChallengeInfo object
    /// containing at least the IV, EncryptedSecret, and Verification fields. These fields are combined with the Yubikey response
    /// to decrypt and verify the secret. 
    /// </summary>
    /// <param name="inf">A pre-populated object containing minimally the IV, EncryptedSecret and Verification fields. 
    ///                   This should be populated from the database.xml auxilliary file</param>
    /// <param name="resp"	>The Yubikey's response to the issued challenge</param>
    /// <returns>The common secret, used as a composite key to encrypt a Keepass database</returns>
    public byte[] GetSecret(ChallengeInfo inf, byte[] resp)
    {
      if (resp.Length != responseLenBytes)
        return null;
      if (inf == null)
        return null;
      if (inf.Challenge == null ||
          inf.Verification == null)
        return null;

      LT64 = inf.LT64;

      byte[] secret;

      if (DecryptSecret(resp, inf, out secret))
      {
        return secret;
      }
      else
      {
        return null;
      }
    }

  }
}
